//! Flight SQL服务实现

use std::pin::Pin;
use std::sync::Arc;
use std::collections::HashMap;
use futures::Stream;
use tonic::{Request, Response, Status, Streaming};
use arrow_flight::{
    flight_service_server::FlightService,
    Action, ActionType, Criteria, Empty, FlightData, FlightDescriptor, 
    FlightInfo, HandshakeRequest, HandshakeResponse, PutResult, 
    SchemaResult, Ticket,
};
use arrow::ipc::writer::IpcWriteOptions;
use arrow::record_batch::RecordBatch;
use data_gateway_kernel::{DataGatewayKernel, Query, QueryType, Result as KernelResult};

/// Flight SQL服务
pub struct FlightSqlService {
    /// 内核引用
    kernel: Arc<DataGatewayKernel>,
    /// 活跃查询
    active_queries: Arc<tokio::sync::RwLock<HashMap<String, QueryHandle>>>,
}

/// 查询句柄
#[derive(Debug, Clone)]
struct QueryHandle {
    /// 查询ID
    query_id: String,
    /// 查询SQL
    sql: String,
    /// 创建时间
    created_at: chrono::DateTime<chrono::Utc>,
    /// 查询状态
    status: QueryStatus,
}

/// 查询状态
#[derive(Debug, Clone)]
enum QueryStatus {
    Pending,
    Executing,
    Completed,
    Failed(String),
}

impl FlightSqlService {
    /// 创建新的Flight SQL服务
    pub fn new(kernel: Arc<DataGatewayKernel>) -> Self {
        Self {
            kernel,
            active_queries: Arc::new(tokio::sync::RwLock::new(HashMap::new())),
        }
    }
    
    /// 执行SQL查询
    async fn execute_sql_query(&self, sql: &str) -> Result<RecordBatch, Status> {
        tracing::debug!("执行SQL查询: {}", sql);
        
        // 创建查询对象
        let query = Query::sql(sql.to_string());
        let query_id = query.id.clone();
        
        // 注册查询句柄
        {
            let mut queries = self.active_queries.write().await;
            queries.insert(query_id.clone(), QueryHandle {
                query_id: query_id.clone(),
                sql: sql.to_string(),
                created_at: chrono::Utc::now(),
                status: QueryStatus::Pending,
            });
        }
        
        // 更新查询状态为执行中
        {
            let mut queries = self.active_queries.write().await;
            if let Some(handle) = queries.get_mut(&query_id) {
                handle.status = QueryStatus::Executing;
            }
        }
        
        // 执行查询
        match self.kernel.handle_query(query).await {
            Ok(result) => {
                // 更新查询状态为完成
                {
                    let mut queries = self.active_queries.write().await;
                    if let Some(handle) = queries.get_mut(&query_id) {
                        handle.status = QueryStatus::Completed;
                    }
                }
                
                // 转换结果为RecordBatch
                if let Some(batches) = result.as_arrow() {
                    if !batches.is_empty() {
                        Ok(batches[0].clone())
                    } else {
                        // 返回空的RecordBatch
                        Ok(RecordBatch::new_empty(Arc::new(arrow::datatypes::Schema::empty())))
                    }
                } else {
                    Err(Status::internal("查询结果格式不支持"))
                }
            },
            Err(e) => {
                // 更新查询状态为失败
                {
                    let mut queries = self.active_queries.write().await;
                    if let Some(handle) = queries.get_mut(&query_id) {
                        handle.status = QueryStatus::Failed(e.to_string());
                    }
                }
                
                tracing::error!("SQL查询执行失败: {}", e);
                Err(Status::internal(format!("查询执行失败: {}", e)))
            }
        }
    }
    
    /// 将RecordBatch转换为FlightData流
    fn record_batch_to_flight_data_stream(
        &self,
        record_batch: RecordBatch,
    ) -> Pin<Box<dyn Stream<Item = Result<FlightData, Status>> + Send>> {
        let schema = record_batch.schema();
        let options = IpcWriteOptions::default();
        
        Box::pin(async_stream::stream! {
            // 首先发送Schema
            match arrow::ipc::writer::IpcDataGenerator::schema_to_bytes(&schema, &options) {
                Ok(schema_bytes) => {
                    yield Ok(FlightData {
                        flight_descriptor: None,
                        data_header: schema_bytes.ipc_message,
                        app_metadata: vec![],
                        data_body: schema_bytes.arrow_data,
                    });
                },
                Err(e) => {
                    yield Err(Status::internal(format!("Schema序列化失败: {}", e)));
                    return;
                }
            }
            
            // 然后发送数据
            let data_gen = arrow::ipc::writer::IpcDataGenerator::default();
            match data_gen.encoded_batch(&record_batch, &options) {
                Ok(encoded_batch) => {
                    yield Ok(FlightData {
                        flight_descriptor: None,
                        data_header: encoded_batch.ipc_message,
                        app_metadata: vec![],
                        data_body: encoded_batch.arrow_data,
                    });
                },
                Err(e) => {
                    yield Err(Status::internal(format!("数据序列化失败: {}", e)));
                }
            }
        })
    }
    
    /// 获取活跃查询列表
    pub async fn get_active_queries(&self) -> Vec<QueryHandle> {
        let queries = self.active_queries.read().await;
        queries.values().cloned().collect()
    }
    
    /// 清理完成的查询
    pub async fn cleanup_completed_queries(&self) {
        let mut queries = self.active_queries.write().await;
        let cutoff_time = chrono::Utc::now() - chrono::Duration::hours(1);
        
        queries.retain(|_, handle| {
            match &handle.status {
                QueryStatus::Completed | QueryStatus::Failed(_) => {
                    handle.created_at > cutoff_time
                },
                _ => true,
            }
        });
    }
}

#[tonic::async_trait]
impl FlightService for FlightSqlService {
    type HandshakeStream = Pin<Box<dyn Stream<Item = Result<HandshakeResponse, Status>> + Send>>;
    type ListFlightsStream = Pin<Box<dyn Stream<Item = Result<FlightInfo, Status>> + Send>>;
    type DoGetStream = Pin<Box<dyn Stream<Item = Result<FlightData, Status>> + Send>>;
    type DoPutStream = Pin<Box<dyn Stream<Item = Result<PutResult, Status>> + Send>>;
    type DoActionStream = Pin<Box<dyn Stream<Item = Result<arrow_flight::Result, Status>> + Send>>;
    type ListActionsStream = Pin<Box<dyn Stream<Item = Result<ActionType, Status>> + Send>>;
    type DoExchangeStream = Pin<Box<dyn Stream<Item = Result<FlightData, Status>> + Send>>;
    
    /// 握手协议
    async fn handshake(
        &self,
        _request: Request<Streaming<HandshakeRequest>>,
    ) -> Result<Response<Self::HandshakeStream>, Status> {
        tracing::debug!("Flight SQL握手请求");
        
        let output = async_stream::stream! {
            yield Ok(HandshakeResponse {
                protocol_version: 0,
                payload: b"Flight SQL Server Ready".to_vec(),
            });
        };
        
        Ok(Response::new(Box::pin(output)))
    }
    
    /// 列出可用的Flight
    async fn list_flights(
        &self,
        _request: Request<Criteria>,
    ) -> Result<Response<Self::ListFlightsStream>, Status> {
        tracing::debug!("列出可用Flight");
        
        let output = async_stream::stream! {
            // 返回空列表，实际实现中可以返回可用的查询或表
        };
        
        Ok(Response::new(Box::pin(output)))
    }
    
    /// 获取Flight信息
    async fn get_flight_info(
        &self,
        request: Request<FlightDescriptor>,
    ) -> Result<Response<FlightInfo>, Status> {
        let descriptor = request.into_inner();
        
        tracing::debug!("获取Flight信息: {:?}", descriptor);
        
        // 解析SQL查询
        if let Some(cmd) = descriptor.cmd.as_ref() {
            let sql = String::from_utf8_lossy(cmd);
            tracing::debug!("解析SQL: {}", sql);
            
            // 执行查询获取Schema信息
            match self.execute_sql_query(&sql).await {
                Ok(record_batch) => {
                    let schema = record_batch.schema();
                    let schema_bytes = match arrow::ipc::writer::IpcDataGenerator::schema_to_bytes(
                        &schema, 
                        &IpcWriteOptions::default()
                    ) {
                        Ok(bytes) => bytes.ipc_message,
                        Err(e) => {
                            return Err(Status::internal(format!("Schema序列化失败: {}", e)));
                        }
                    };
                    
                    let flight_info = FlightInfo {
                        schema: schema_bytes,
                        flight_descriptor: Some(descriptor),
                        endpoint: vec![],
                        total_records: record_batch.num_rows() as i64,
                        total_bytes: -1, // 未知大小
                        ordered: false,
                    };
                    
                    Ok(Response::new(flight_info))
                },
                Err(status) => Err(status),
            }
        } else {
            Err(Status::invalid_argument("缺少SQL命令"))
        }
    }
    
    /// 获取Schema
    async fn get_schema(
        &self,
        request: Request<FlightDescriptor>,
    ) -> Result<Response<SchemaResult>, Status> {
        let descriptor = request.into_inner();
        
        tracing::debug!("获取Schema: {:?}", descriptor);
        
        if let Some(cmd) = descriptor.cmd.as_ref() {
            let sql = String::from_utf8_lossy(cmd);
            
            match self.execute_sql_query(&sql).await {
                Ok(record_batch) => {
                    let schema = record_batch.schema();
                    let schema_bytes = match arrow::ipc::writer::IpcDataGenerator::schema_to_bytes(
                        &schema,
                        &IpcWriteOptions::default()
                    ) {
                        Ok(bytes) => bytes.ipc_message,
                        Err(e) => {
                            return Err(Status::internal(format!("Schema序列化失败: {}", e)));
                        }
                    };
                    
                    let schema_result = SchemaResult { schema: schema_bytes };
                    Ok(Response::new(schema_result))
                },
                Err(status) => Err(status),
            }
        } else {
            Err(Status::invalid_argument("缺少SQL命令"))
        }
    }
    
    /// 执行查询并返回数据流
    async fn do_get(
        &self,
        request: Request<Ticket>,
    ) -> Result<Response<Self::DoGetStream>, Status> {
        let ticket = request.into_inner();
        
        tracing::debug!("执行DoGet: {:?}", ticket);
        
        let sql = String::from_utf8_lossy(&ticket.ticket);
        
        match self.execute_sql_query(&sql).await {
            Ok(record_batch) => {
                let stream = self.record_batch_to_flight_data_stream(record_batch);
                Ok(Response::new(stream))
            },
            Err(status) => Err(status),
        }
    }
    
    /// 上传数据（暂不实现）
    async fn do_put(
        &self,
        _request: Request<Streaming<FlightData>>,
    ) -> Result<Response<Self::DoPutStream>, Status> {
        Err(Status::unimplemented("DoPut暂不支持"))
    }
    
    /// 执行动作
    async fn do_action(
        &self,
        request: Request<Action>,
    ) -> Result<Response<Self::DoActionStream>, Status> {
        let action = request.into_inner();
        
        tracing::debug!("执行Action: {}", action.r#type);
        
        let output = async_stream::stream! {
            match action.r#type.as_str() {
                "get_active_queries" => {
                    let queries = self.get_active_queries().await;
                    let result = serde_json::to_vec(&queries).unwrap_or_default();
                    yield Ok(arrow_flight::Result { body: result });
                },
                "cleanup_queries" => {
                    self.cleanup_completed_queries().await;
                    yield Ok(arrow_flight::Result { body: b"OK".to_vec() });
                },
                _ => {
                    yield Err(Status::unimplemented(format!("不支持的Action: {}", action.r#type)));
                }
            }
        };
        
        Ok(Response::new(Box::pin(output)))
    }
    
    /// 列出可用动作
    async fn list_actions(
        &self,
        _request: Request<Empty>,
    ) -> Result<Response<Self::ListActionsStream>, Status> {
        tracing::debug!("列出可用Action");
        
        let output = async_stream::stream! {
            yield Ok(ActionType {
                r#type: "get_active_queries".to_string(),
                description: "获取活跃查询列表".to_string(),
            });
            
            yield Ok(ActionType {
                r#type: "cleanup_queries".to_string(),
                description: "清理完成的查询".to_string(),
            });
        };
        
        Ok(Response::new(Box::pin(output)))
    }
    
    /// 双向流交换（暂不实现）
    async fn do_exchange(
        &self,
        _request: Request<Streaming<FlightData>>,
    ) -> Result<Response<Self::DoExchangeStream>, Status> {
        Err(Status::unimplemented("DoExchange暂不支持"))
    }
}
